#!/usr/bin/env python3
"""
告警管理模块
负责处理和发送系统资源告警
"""

import os
import json
import time
import logging
import smtplib
import requests
from typing import Dict, List, Any, Optional
from datetime import datetime
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart

# 配置日志
logger = logging.getLogger('AlertManager')
logger.setLevel(logging.INFO)

class AlertManager:
    """告警管理类"""
    def __init__(self, config_path: str = 'config/alert_config.json'):
        # 使用绝对路径处理配置文件
        base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        self.config_path = os.path.join(base_dir, config_path)
        self.config = self._load_config(self.config_path)
        self.alert_history: List[Dict[str, Any]] = []
        self.alert_cooldown = 300  # 告警冷却时间（秒）
    
    def _load_config(self, config_path: str) -> Dict[str, Any]:
        """加载告警配置"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            logger.error(f"加载告警配置失败: {str(e)}")
            # 返回默认配置
            default_config = {
                'email': {
                    'enabled': False,
                    'smtp_server': 'smtp.example.com',
                    'smtp_port': 587,
                    'username': '',
                    'password': '',
                    'from_addr': '',
                    'to_addrs': []
                },
                'dingtalk': {
                    'enabled': False,
                    'webhook_url': '',
                    'secret': ''
                },
                'notification_levels': ['critical']
            }
            
            # 尝试创建默认配置文件
            try:
                os.makedirs(os.path.dirname(config_path), exist_ok=True)
                with open(config_path, 'w', encoding='utf-8') as f:
                    json.dump(default_config, f, ensure_ascii=False, indent=4)
                logger.info(f"已创建默认告警配置文件: {config_path}")
            except Exception as e:
                logger.error(f"创建默认告警配置文件失败: {str(e)}")
                
            return default_config
    
    def process_alert(self, alert_data: Dict[str, Any]) -> None:
        """处理告警信息"""
        if not self._should_send_alert(alert_data):
            return
        
        # 格式化告警消息
        message = self._format_alert_message(alert_data)
        
        # 根据告警级别发送通知
        if alert_data['status'] in self.config['notification_levels']:
            self._send_notifications(message, alert_data)
        
        # 记录告警历史
        self._record_alert(alert_data)
    
    def _should_send_alert(self, alert_data: Dict[str, Any]) -> bool:
        """判断是否应该发送告警"""
        if alert_data['status'] == 'normal':
            return False
        
        # 检查是否在冷却期内
        now = time.time()
        for alert in reversed(self.alert_history):
            if alert['status'] == alert_data['status'] and \
               now - alert['timestamp'] < self.alert_cooldown:
                return False
        
        return True
    
    def _format_alert_message(self, alert_data: Dict[str, Any]) -> str:
        """格式化告警消息"""
        status_map = {'critical': '严重', 'warning': '警告', 'normal': '正常'}
        status = status_map.get(alert_data['status'], alert_data['status'])
        
        message = f"系统资源告警 - {status}\n\n"
        message += f"时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
        message += "告警详情:\n"
        
        for warning in alert_data.get('warnings', []):
            message += f"- {warning}\n"
        
        if 'metrics' in alert_data:
            metrics = alert_data['metrics']
            message += "\n系统指标:\n"
            message += f"- CPU使用率: {metrics['cpu']['percent']}%\n"
            message += f"- 内存使用率: {metrics['memory']['percent']}%\n"
            message += f"- 磁盘使用率: {metrics['disk']['percent']}%\n"
        
        return message
    
    def _send_notifications(self, message: str, alert_data: Dict[str, Any]) -> None:
        """发送通知"""
        # 发送邮件通知
        if self.config['email']['enabled']:
            self._send_email(message)
        
        # 发送钉钉通知
        if self.config['dingtalk']['enabled']:
            self._send_dingtalk(message)
    
    def _send_email(self, message: str) -> None:
        """发送邮件通知"""
        try:
            email_config = self.config['email']
            msg = MIMEMultipart()
            msg['From'] = email_config['from_addr']
            msg['To'] = ', '.join(email_config['to_addrs'])
            msg['Subject'] = '系统资源告警'
            
            msg.attach(MIMEText(message, 'plain', 'utf-8'))
            
            with smtplib.SMTP(email_config['smtp_server'], email_config['smtp_port']) as server:
                server.starttls()
                server.login(email_config['username'], email_config['password'])
                server.send_message(msg)
            
            logger.info("邮件告警发送成功")
        except Exception as e:
            logger.error(f"发送邮件告警失败: {str(e)}")
    
    def _send_dingtalk(self, message: str) -> None:
        """发送钉钉通知"""
        try:
            dingtalk_config = self.config['dingtalk']
            headers = {'Content-Type': 'application/json'}
            
            data = {
                'msgtype': 'text',
                'text': {
                    'content': message
                }
            }
            
            response = requests.post(
                dingtalk_config['webhook_url'],
                headers=headers,
                json=data
            )
            
            if response.status_code == 200:
                logger.info("钉钉告警发送成功")
            else:
                logger.error(f"钉钉告警发送失败: {response.text}")
        except Exception as e:
            logger.error(f"发送钉钉告警失败: {str(e)}")
    
    def _record_alert(self, alert_data: Dict[str, Any]) -> None:
        """记录告警历史"""
        alert_record = {
            'timestamp': time.time(),
            'datetime': datetime.now().isoformat(),
            'status': alert_data['status'],
            'warnings': alert_data.get('warnings', []),
            'metrics': alert_data.get('metrics', {})
        }
        
        self.alert_history.append(alert_record)
        
        # 只保留最近100条记录
        if len(self.alert_history) > 100:
            self.alert_history = self.alert_history[-100:]
    
    def get_alert_history(self, limit: int = 10) -> List[Dict[str, Any]]:
        """获取告警历史记录"""
        return self.alert_history[-limit:]
    
    def get_config(self) -> Dict[str, Any]:
        """获取当前告警配置"""
        return self.config
    
    def update_config(self, new_config: Dict[str, Any]) -> None:
        """更新告警配置"""
        try:
            # 更新内存中的配置
            self.config.update(new_config)
            
            # 保存到文件
            os.makedirs(os.path.dirname(self.config_path), exist_ok=True)
            
            with open(self.config_path, 'w', encoding='utf-8') as f:
                json.dump(self.config, f, ensure_ascii=False, indent=4)
                
            logger.info(f"告警配置已更新并保存到: {self.config_path}")
        except Exception as e:
            logger.error(f"更新告警配置失败: {str(e)}")
            raise